#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
交互式频散曲线拾取程序
支持从FV图中人工拾取频散曲线，并自动调整到局部极大值点
"""

import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, TextBox
import h5py
import os
from datetime import datetime

from utils.h5data import read_h5_data
from utils.math import norm

# 配置参数
CONFIG = {
    'DPI': 100,
    'FIGSIZE': (15, 10),
    'SEARCH_RANGE': 50,  # 速度搜索范围（米/秒）
    'DOUBLE_CLICK_SEARCH_RANGE': 20,  # 双击时的速度搜索范围（米/秒）
    'FV_NORM': True,
    'DEFAULT_SCALE': 60
}

def do_fv(data, vs, x=None, sx=0, dt=1e-2, NORM=False, PRINT=False, fs=-1, fe=1e10):
    """执行频散变换"""
    I = np.sqrt(np.array(-1, dtype='complex'))
    PI = np.pi

    nx, nt = data.shape
    if x is None:
        x = np.arange(nx) * 50  # 默认间距50米
    t = np.arange(nt) * dt

    XF = np.zeros([nx, nt], dtype='complex')
    for ix in range(nx):
        XF[ix, :] = np.fft.fft(data[ix, :])
        if PRINT:
            print('\r FX:{}|{}'.format(ix, nx), end='', flush=True)
    FX = XF.T
    f = np.arange(nt) * 1 / t.max()
    if PRINT:
        print('')
    nv = len(vs)
    FV = np.zeros([nt, nv], dtype='complex')
    idx_valid = []
    for j in range(nt // 2):
        if f[j] < fs or f[j] > fe:
            continue
        idx_valid.append(j)
        for i, v in enumerate(vs):
            # FV[j,i] = np.sum(np.exp(I*2*PI*np.abs(x-sx)/v*f[j])*FX[j,:])
            FV[j, i] = np.sum(np.exp(I * 2 * PI * (x - sx) / v * f[j]) * FX[j, :])
        if NORM:
            FV[j, :] = FV[j, :] / np.abs(FV[j, :]).max()
        if PRINT:
            print('\r FV:{}|{}'.format(j, nt), end='')

    return FV[idx_valid, :], f[idx_valid], vs


class DispersionCurvePicker:
    def __init__(self, h5_file):
        """初始化频散曲线拾取器"""
        self.h5_file = h5_file
        
        # 初始化默认参数
        self.vs = np.linspace(100, 1400, 800)  # 默认速度范围
        self.fs = 1.0  # 默认频率下限
        self.fe = 5.0  # 默认频率上限
        self.SCALE = CONFIG['DEFAULT_SCALE']  # 波形显示的SCALE初始值
        
        # 读取HDF5文件中的数据
        with h5py.File(h5_file, 'r') as f:
            # 尝试读取FV相关数据
            if 'FV' in f and 'vs' in f and 'f' in f:
                self.FV = np.abs(f['FV'][()])  # 取模值
                self.vs = f['vs'][()]
                self.f = f['f'][()]
                self.have_fv_data = True
            else:
                self.have_fv_data = False
                
            # 读取波形数据（必须存在）
            self.traces = f['traces'][()]
            self.x = f['x'][()]
            self.t = f['t'][()]
            
        # 如果没有FV数据，则计算默认的FV矩阵
        if not self.have_fv_data:
            self.log_action("FV data not found", "Computing FV matrix with default parameters")
            self.dt = self.t[1] - self.t[0] if len(self.t) > 1 else 1e-2
            FV, f_valid, vs_valid = do_fv(self.traces, self.vs, x=np.array(self.x), 
                                          sx=0, dt=self.dt, NORM=CONFIG['FV_NORM'], 
                                          PRINT=False, fs=self.fs, fe=self.fe)
            self.FV = np.abs(FV)
            self.f = f_valid
            self.vs = vs_valid
        
        # 归一化FV图
        self.FV = self.FV / np.max(self.FV)
        
        # 保存原始数据范围
        self.original_vs = self.vs.copy()
        self.original_f = self.f.copy()
        self.original_FV = self.FV.copy()
        
        # 初始化变量
        self.nv = len(self.vs)
        self.nt = len(self.f)
        self.current_curve = None  # 当前拾取的曲线
        self.current_curve_index = 0  # 当前曲线索引（支持多条曲线）
        self.curves = {}  # 存储所有拾取的曲线
        self.selected_point = None  # 当前选中的点
        self.mode = 'add'  # 模式：add（添加点）或 move（移动点）
        self.selected_point_marker = None  # 用于标记选中点的图形对象
        self.status_text = None  # 用于显示当前状态的文本
        self.dt = self.t[1] - self.t[0] if len(self.t) > 1 else 1e-2  # 时间采样间隔
        self.colorbar = None  # 用于存储colorbar对象
        
        # 初始化日志
        self.log_file = self.create_log_file()
        
        # 设置图形界面（如果尚未设置）
        if not hasattr(self, 'fig') or not plt.fignum_exists(self.fig.number):
            self.setup_figure()
        else:
            # 更新FV显示
            self.update_fv_display()
            # 清除曲线
            self.clear_all_curves()
            
        # 基于能量最大值初始化参考曲线
        self.initialize_reference_curve()
        
    def setup_figure(self):
        """设置图形界面"""
        self.fig = plt.figure(figsize=CONFIG['FIGSIZE'], dpi=CONFIG['DPI'])
        
        # 创建子图
        self.ax_fv = plt.subplot(2, 2, 1)
        self.ax_waveform = plt.subplot(2, 2, 2)
        self.ax_controls = plt.subplot(2, 2, 3)  # 将ax_controls移到subplot(223)
        self.ax_trace = plt.subplot(2, 2, 4)
        
        # 隐藏控制子图的坐标轴
        self.ax_controls.axis('off')
        
        # 显示FV图 (f为横轴，v为纵轴，速度轴上大下小)
        extent = [self.f[0], self.f[-1], self.vs[-1], self.vs[0]]
        self.im = self.ax_fv.imshow(self.FV.T, extent=extent, aspect='auto', cmap='jet')
        self.ax_fv.set_xlabel('Frequency (Hz)')
        self.ax_fv.set_ylabel('Phase Velocity (m/s)')
        self.ax_fv.set_title('Frequency-Velocity Spectrum')
        # 不再需要invert_yaxis()，因为extent已经正确设置了速度轴方向
        
        # 显示波形数据
        if len(self.traces) > 0:
            self.plot_waveforms()
        
        # 初始化曲线绘制
        self.curve_lines = {}
        self.point_annotations = {}
        
        # 连接鼠标事件
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)
        
        # 添加按钮和输入框
        self.add_controls()
        
        # 添加状态显示文本
        self.status_text = self.ax_controls.text(0.1, 0.9, '', 
                                                transform=self.ax_controls.transAxes,
                                                verticalalignment='top',
                                                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        self.update_status_text()
        
        # 添加说明文本
        self.ax_controls.text(0.1, 0.6, 'Controls:\n' +
                             'Left click: Add/move point\n' +
                             'Double click: Force select point\n' +
                             'Right click: Remove point\n' +
                             'Keys:\n' +
                             '  a: Add mode\n' +
                             '  m: Move mode\n' +
                             '  r: Remove mode\n' +
                             '  s: Save curves\n' +
                             '  n: New curve\n' +
                             '  d: Delete current curve\n' +
                             '  l: Load curves from file\n' +
                             '  del: Delete selected point\n' +
                             'Text Inputs:\n' +
                             '  Ref: Reference .disp file\n' +
                             '  H5: H5 file path\n',
                             transform=self.ax_controls.transAxes,
                             verticalalignment='top',
                             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
    def update_status_text(self):
        """更新状态显示文本"""
        if self.status_text:
            self.status_text.set_text(f'Mode: {self.mode.upper()}\nCurve: {self.current_curve_index}\n' +
                                     f'V: [{self.vs[0]:.0f}, {self.vs[-1]:.0f}] m/s\n' +
                                     f'F: [{self.f[0]:.1f}, {self.f[-1]:.1f}] Hz\n' +
                                     f'Scale: {self.SCALE}')
        
    def plot_waveforms(self):
        """绘制所有波形数据"""
        # 清除之前的波形图
        self.ax_waveform.clear()
        
        nx = len(self.traces)
        if nx == 0:
            return
            
        for i in range(nx):
            trace_i = self.traces[i, :].copy()
            trace_i = norm(trace_i)
            trace_i = trace_i * self.SCALE
            self.ax_waveform.plot(self.t, trace_i + self.x[i], lw=0.5, color='k')
            
        self.ax_waveform.set_xlabel('Time (s)')
        self.ax_waveform.set_ylabel('Distance (m)')
        self.ax_waveform.set_title('Waveforms')
        self.fig.canvas.draw()
        
    def add_controls(self):
        """添加控制按钮和输入框"""
        # 按钮和输入框位置
        button_width = 0.12
        button_height = 0.04
        button_spacing = 0.01
        input_width = 0.08
        input_height = 0.04
        
        # 第一列按钮
        ax_add = plt.axes([0.6, 0.45, button_width, button_height])
        self.btn_add = Button(ax_add, 'Add Mode')
        self.btn_add.on_clicked(self.set_add_mode)
        
        ax_move = plt.axes([0.6, 0.45 - (button_height + button_spacing), button_width, button_height])
        self.btn_move = Button(ax_move, 'Move Mode')
        self.btn_move.on_clicked(self.set_move_mode)
        
        ax_remove = plt.axes([0.6, 0.45 - 2*(button_height + button_spacing), button_width, button_height])
        self.btn_remove = Button(ax_remove, 'Remove Mode')
        self.btn_remove.on_clicked(self.set_remove_mode)
        
        ax_new = plt.axes([0.6, 0.45 - 3*(button_height + button_spacing), button_width, button_height])
        self.btn_new = Button(ax_new, 'New Curve')
        self.btn_new.on_clicked(self.new_curve)
        
        ax_delete = plt.axes([0.6, 0.45 - 4*(button_height + button_spacing), button_width, button_height])
        self.btn_delete = Button(ax_delete, 'Delete Curve')
        self.btn_delete.on_clicked(self.delete_current_curve)
        
        ax_save = plt.axes([0.6, 0.45 - 5*(button_height + button_spacing), button_width, button_height])
        self.btn_save = Button(ax_save, 'Save Curves')
        self.btn_save.on_clicked(self.save_curves)
        
        ax_load = plt.axes([0.6, 0.45 - 6*(button_height + button_spacing), button_width, button_height])
        self.btn_load = Button(ax_load, 'Load Curves')
        self.btn_load.on_clicked(self.load_curves_from_file)
        
        # 第二列按钮和输入框
        # 速度范围输入
        ax_vmin_text = plt.axes([0.8, 0.45, 0.04, button_height])
        ax_vmin_text.text(0, 0, 'V min:', transform=ax_vmin_text.transAxes, verticalalignment='center')
        ax_vmin_text.axis('off')
        
        ax_vmin = plt.axes([0.85, 0.45, input_width, input_height])
        self.txt_vmin = TextBox(ax_vmin, '', initial=str(int(self.vs[0])))
        self.txt_vmin.on_submit(self.on_text_change)
        
        ax_vmax_text = plt.axes([0.8, 0.45 - (button_height + button_spacing), 0.04, button_height])
        ax_vmax_text.text(0, 0, 'V max:', transform=ax_vmax_text.transAxes, verticalalignment='center')
        ax_vmax_text.axis('off')
        
        ax_vmax = plt.axes([0.85, 0.45 - (button_height + button_spacing), input_width, input_height])
        self.txt_vmax = TextBox(ax_vmax, '', initial=str(int(self.vs[-1])))
        self.txt_vmax.on_submit(self.on_text_change)
        
        # 频率范围输入
        ax_fmin_text = plt.axes([0.8, 0.45 - 2*(button_height + button_spacing), 0.04, button_height])
        ax_fmin_text.text(0, 0, 'F min:', transform=ax_fmin_text.transAxes, verticalalignment='center')
        ax_fmin_text.axis('off')
        
        ax_fmin = plt.axes([0.85, 0.45 - 2*(button_height + button_spacing), input_width, input_height])
        self.txt_fmin = TextBox(ax_fmin, '', initial=f"{self.f[0]:.1f}")
        self.txt_fmin.on_submit(self.on_text_change)
        
        ax_fmax_text = plt.axes([0.8, 0.45 - 3*(button_height + button_spacing), 0.04, button_height])
        ax_fmax_text.text(0, 0, 'F max:', transform=ax_fmax_text.transAxes, verticalalignment='center')
        ax_fmax_text.axis('off')
        
        ax_fmax = plt.axes([0.85, 0.45 - 3*(button_height + button_spacing), input_width, input_height])
        self.txt_fmax = TextBox(ax_fmax, '', initial=f"{self.f[-1]:.1f}")
        self.txt_fmax.on_submit(self.on_text_change)
        
        # SCALE输入
        ax_scale_text = plt.axes([0.8, 0.45 - 4*(button_height + button_spacing), 0.04, button_height])
        ax_scale_text.text(0, 0, 'Scale:', transform=ax_scale_text.transAxes, verticalalignment='center')
        ax_scale_text.axis('off')
        
        ax_scale = plt.axes([0.85, 0.45 - 4*(button_height + button_spacing), input_width, input_height])
        self.txt_scale = TextBox(ax_scale, '', initial=str(self.SCALE))
        self.txt_scale.on_submit(self.update_scale)
        
        # 更新FV按钮
        ax_update_fv = plt.axes([0.8, 0.45 - 5*(button_height + button_spacing), button_width, button_height])
        self.btn_update_fv = Button(ax_update_fv, 'Update FV')
        self.btn_update_fv.on_clicked(self.update_fv_matrix)
        
        # 参考曲线文件输入
        ax_ref_text = plt.axes([0.8, 0.45 - 6*(button_height + button_spacing), 0.04, button_height])
        ax_ref_text.text(0, 0, 'Ref:', transform=ax_ref_text.transAxes, verticalalignment='center')
        ax_ref_text.axis('off')
        
        ax_ref = plt.axes([0.85, 0.45 - 6*(button_height + button_spacing), input_width, input_height])
        self.txt_ref = TextBox(ax_ref, '', initial='')
        self.txt_ref.on_submit(self.load_reference_curves)
        
        # H5文件输入和刷新按钮
        ax_h5_text = plt.axes([0.05, 0.25, 0.04, button_height])
        ax_h5_text.text(0, 0, 'H5:', transform=ax_h5_text.transAxes, verticalalignment='center')
        ax_h5_text.axis('off')
        
        ax_h5 = plt.axes([0.1, 0.25, 0.25, input_height])
        self.txt_h5 = TextBox(ax_h5, '', initial=self.h5_file if hasattr(self, 'h5_file') else '')
        
        ax_refresh = plt.axes([0.37, 0.25, 0.08, button_height])
        self.btn_refresh = Button(ax_refresh, 'Refresh')
        self.btn_refresh.on_clicked(self.refresh_h5_file)
        
    def on_text_change(self, text):
        """处理文本框内容变化，但不立即更新FV，只更新scale"""
        try:
            scale = float(self.txt_scale.text)
            if scale > 0:
                self.SCALE = scale
                self.plot_waveforms()
                self.update_status_text()
        except ValueError:
            pass  # 忽略无效输入
            
    def update_scale(self, text):
        """更新SCALE值"""
        try:
            scale = float(text)
            if scale > 0:
                self.SCALE = scale
                self.plot_waveforms()
                self.update_status_text()
        except ValueError:
            pass  # 忽略无效输入
            
    def update_fv_matrix(self, event):
        """手动更新FV矩阵"""
        try:
            vmin = float(self.txt_vmin.text)
            vmax = float(self.txt_vmax.text)
            fmin = float(self.txt_fmin.text)
            fmax = float(self.txt_fmax.text)
            scale = float(self.txt_scale.text)
            
            # 记录参数更新
            self.log_action("Parameters updated", f"V: [{vmin}, {vmax}] m/s, F: [{fmin}, {fmax}] Hz, Scale: {scale}")
            
            # 更新SCALE
            if scale > 0:
                self.SCALE = scale
                self.plot_waveforms()
            
            # 重新计算FV矩阵
            self.recalculate_fv(vmin=vmin, vmax=vmax, fmin=fmin, fmax=fmax)
        except ValueError:
            self.log_action("Error", "Invalid input values for FV matrix update")
            
    def recalculate_fv(self, vmin=None, vmax=None, fmin=None, fmax=None):
        """重新计算FV矩阵"""
        # 获取当前范围
        current_vmin = vmin if vmin is not None else self.vs[0]
        current_vmax = vmax if vmax is not None else self.vs[-1]
        current_fmin = fmin if fmin is not None else self.f[0]
        current_fmax = fmax if fmax is not None else self.f[-1]
        
        # 更新vs和f范围
        if vmin is not None or vmax is not None:
            # 保持800个速度点
            self.vs = np.linspace(current_vmin, current_vmax, 800)
            
        if fmin is not None or fmax is not None:
            # 频率点根据原始数据确定
            valid_freq_indices = np.where((self.original_f >= current_fmin) & (self.original_f <= current_fmax))[0]
            if len(valid_freq_indices) > 0:
                self.f = self.original_f[valid_freq_indices]
            else:
                # 如果没有有效频率点，创建新的
                self.f = np.linspace(current_fmin, current_fmax, 100)  # 默认100个点
                
        # 重新计算FV矩阵
        try:
            FV, f_valid, vs_valid = do_fv(self.traces, self.vs, x=np.array(self.x), 
                                          sx=0, dt=self.dt, NORM=CONFIG['FV_NORM'], PRINT=False, 
                                          fs=current_fmin, fe=current_fmax)
            self.FV = np.abs(FV)
            self.FV = self.FV / np.max(self.FV) if np.max(self.FV) > 0 else self.FV
            
            # 更新显示
            self.update_fv_display()
            
            # 清空所有曲线并生成新的参考曲线
            self.clear_all_curves()
            self.initialize_reference_curve()
            
            # 更新状态显示
            self.update_status_text()
            
            self.log_action("FV matrix recalculated", f"V: [{current_vmin}, {current_vmax}] m/s, F: [{current_fmin}, {current_fmax}] Hz")
        except Exception as e:
            self.log_action("Error recalculating FV matrix", str(e))
            
    def update_fv_display(self):
        """更新FV图像显示"""
        # 清除之前的图像
        self.ax_fv.clear()
        
        # 显示新的FV图 (f为横轴，v为纵轴，速度轴上大下小)
        extent = [self.f[0], self.f[-1], self.vs[-1], self.vs[0]]
        self.im = self.ax_fv.imshow(self.FV.T, extent=extent, aspect='auto', cmap='jet')
        self.ax_fv.set_xlabel('Frequency (Hz)')
        self.ax_fv.set_ylabel('Phase Velocity (m/s)')
        self.ax_fv.set_title('Frequency-Velocity Spectrum')
        
        self.fig.canvas.draw()
        
    def clear_all_curves(self):
        """清空所有曲线"""
        # 删除所有曲线的可视化元素
        for curve_index in list(self.curve_lines.keys()):
            try:
                self.curve_lines[curve_index].remove()
            except NotImplementedError:
                # 如果无法直接移除，将其设为不可见
                self.curve_lines[curve_index].set_visible(False)
            # 清除标注
            for ann in self.point_annotations.get(curve_index, []):
                try:
                    ann.remove()
                except NotImplementedError:
                    # 如果无法直接移除，将其设为不可见
                    ann.set_visible(False)
        
        # 清空存储
        self.curve_lines = {}
        self.point_annotations = {}
        self.curves = {}
        self.current_curve = None
        self.current_curve_index = 0
        
    def initialize_reference_curve(self):
        """基于每个频率点上的最大能量初始化参考曲线"""
        # 首先尝试加载参考曲线文件
        reference_file = None
        if hasattr(self, 'reference_file') and self.reference_file:
            reference_file = self.reference_file
        else:
            # 尝试加载与h5文件同名的disp文件
            disp_file = os.path.splitext(self.h5_file)[0] + '.disp'
            if os.path.exists(disp_file):
                reference_file = disp_file
                
        if reference_file and os.path.exists(reference_file):
            self.load_dispersion_curves(reference_file)
            self.log_action("Loaded reference dispersion curves", f"From {reference_file}")
        else:
            # 如果没有参考文件，则基于能量最大值初始化
            if len(self.FV) == 0 or len(self.f) == 0 or len(self.vs) == 0:
                return
                
            curve_points = []
            
            # 对每个频率找到最大能量对应的速度
            for i in range(len(self.f)):
                # 找到该频率下能量最大的速度点
                max_idx = np.argmax(self.FV[i, :])
                curve_points.append((self.f[i], self.vs[max_idx]))
            
            # 保存为第一条曲线
            self.curves[0] = np.array(curve_points)
            self.current_curve_index = 0
            self.current_curve = self.curves[0]
            self.draw_curve()
        
    def find_local_maximum(self, freq_val, initial_vs, is_double_click=False):
        """
        在给定频率下，找到initial_vs附近的局部能量极大值点
        支持对已有曲线点进行微调
        """
        # 找到最接近的频率索引
        freq_idx = np.argmin(np.abs(self.f - freq_val))
        
        # 找到最接近的vs索引
        vs_idx = np.argmin(np.abs(self.vs - initial_vs))
        
        # 在该频率下寻找局部极大值
        fv_row = self.FV[freq_idx, :]
        
        # 根据是否是双击选择不同的搜索范围
        if is_double_click:
            search_window_m_s = CONFIG['DOUBLE_CLICK_SEARCH_RANGE']  # 双击时的搜索范围
        else:
            search_window_m_s = CONFIG['SEARCH_RANGE']  # 普通点击的搜索范围
            
        # 计算速度分辨率
        dv = (self.vs[-1] - self.vs[0]) / len(self.vs) if len(self.vs) > 1 else 1
        # 计算搜索窗口点数
        search_window = int(search_window_m_s / dv) if dv > 0 else 100
        search_window = max(1, search_window)  # 至少为1
        
        # 确保搜索窗口在有效范围内
        start_idx = max(0, vs_idx - search_window)
        end_idx = min(len(self.vs), vs_idx + search_window)
        
        # 在指定窗口内找到局部最大值
        local_max_idx = start_idx + np.argmax(fv_row[start_idx:end_idx])
        
        # 返回找到的最大值对应的速度
        return self.vs[local_max_idx]
        
    def draw_curve(self):
        """绘制当前曲线"""
        # 定义颜色列表，为不同曲线使用不同颜色
        colors = ['white', 'red', 'yellow', 'cyan', 'magenta', 'lime', 'orange', 'pink']
        
        # 清除之前的曲线和标注
        if self.current_curve_index in self.curve_lines:
            try:
                self.curve_lines[self.current_curve_index].remove()
            except NotImplementedError:
                # 如果无法直接移除，将其设为不可见
                self.curve_lines[self.current_curve_index].set_visible(False)
            # 同时清除标注
            for ann in self.point_annotations.get(self.current_curve_index, []):
                try:
                    ann.remove()
                except NotImplementedError:
                    # 如果无法直接移除，将其设为不可见
                    ann.set_visible(False)
            if self.current_curve_index in self.point_annotations:
                del self.point_annotations[self.current_curve_index]
            
        # 绘制当前曲线
        if self.current_curve is not None and len(self.current_curve) > 0:
            # 根据曲线索引选择颜色
            color_idx = self.current_curve_index % len(colors)
            curve_color = colors[color_idx]
            
            line, = self.ax_fv.plot(self.current_curve[:, 0], self.current_curve[:, 1], 
                                   'o-', color=curve_color, linewidth=2, markersize=4,
                                   label=f'Curve {self.current_curve_index}', picker=True)
            self.curve_lines[self.current_curve_index] = line
            
            # 添加标注
            self.point_annotations[self.current_curve_index] = []
            
            for i, (freq, vs) in enumerate(self.current_curve):
                ann = self.ax_fv.annotate(str(i), (freq, vs), 
                                         xytext=(5, 5), textcoords='offset points',
                                         color=curve_color, fontsize=8)
                self.point_annotations[self.current_curve_index].append(ann)
        
        self.ax_fv.legend()
        self.fig.canvas.draw()
        
    def on_click(self, event):
        """处理鼠标点击事件"""
        # 只处理FV图上的点击
        if event.inaxes != self.ax_fv:
            return
            
        freq = event.xdata
        vs = event.ydata
        
        if freq is None or vs is None:
            return
            
        # 检查是否是双击
        if event.dblclick:
            # 双击强制选择点，在更小范围内搜索最大值
            adjusted_vs = self.find_local_maximum(freq, vs, is_double_click=True)
            self.force_select_point(freq, adjusted_vs)
            return
            
        # 找到最近的频率索引
        freq_idx = np.argmin(np.abs(self.f - freq))
        nearest_freq = self.f[freq_idx]
        
        if event.button == 1:  # 左键 - 添加或移动点
            if self.mode == 'add':
                self.add_point(freq_idx, nearest_freq, vs, is_double_click=False)
            elif self.mode == 'move':
                self.move_point(freq_idx, nearest_freq, vs, is_double_click=False)
            # 标记选中的点为红色
            self.mark_selected_point(freq, vs)
                
        elif event.button == 3:  # 右键 - 删除点（在所有模式下都可用）
            self.remove_point(freq_idx, vs)
                
    def force_select_point(self, freq, vs):
        """双击强制选择点"""
        self.log_action("Point force selected", f"f={freq:.2f} Hz, v={vs:.0f} m/s")
        # 标记选中的点为红色
        self.mark_selected_point(freq, vs)
                
    def mark_selected_point(self, freq, vs):
        """标记选中的点为红色"""
        # 清除之前的标记
        if self.selected_point_marker:
            try:
                self.selected_point_marker.remove()
            except NotImplementedError:
                self.selected_point_marker.set_visible(False)
            
        # 添加新的标记
        self.selected_point_marker, = self.ax_fv.plot(freq, vs, 'ro', markersize=8, markeredgewidth=2)
        self.fig.canvas.draw()
                
    def add_point(self, freq_idx, freq, vs, is_double_click=False):
        """添加点到曲线上"""
        if self.current_curve is None:
            self.current_curve = np.empty((0, 2))
            
        # 调整到局部极大值
        adjusted_vs = self.find_local_maximum(freq, vs, is_double_click=is_double_click)
        
        # 检查是否已经有点在这个频率上
        freq_exists = False
        point_index = -1
        for i, (existing_freq, _) in enumerate(self.current_curve):
            if np.isclose(existing_freq, freq, atol=0.01):
                # 更新现有频率点的速度
                old_vs = self.current_curve[i][1]
                self.current_curve[i] = [freq, adjusted_vs]
                freq_exists = True
                point_index = i
                self.log_action("Point moved", f"Curve {self.current_curve_index}, f={freq:.2f} Hz, v: {old_vs:.0f} -> {adjusted_vs:.0f} m/s")
                break
                
        if not freq_exists:
            # 添加新点
            new_point = np.array([[freq, adjusted_vs]])
            self.current_curve = np.vstack([self.current_curve, new_point]) if len(self.current_curve) > 0 else new_point
            point_index = len(self.current_curve) - 1 if len(self.current_curve) > 0 else 0
            self.log_action("Point added", f"Curve {self.current_curve_index}, f={freq:.2f} Hz, v={adjusted_vs:.0f} m/s")
            
        # 按频率排序
        sorted_indices = np.argsort(self.current_curve[:, 0])
        self.current_curve = self.current_curve[sorted_indices]
        
        # 更新存储
        self.curves[self.current_curve_index] = self.current_curve
        self.draw_curve()
        
    def move_point(self, freq_idx, freq, vs, is_double_click=False):
        """移动选中的点"""
        # 这里简化处理，实际应该找到最近的点进行移动
        self.add_point(freq_idx, freq, vs, is_double_click=is_double_click)
        
    def remove_point(self, freq_idx, vs):
        """从曲线上删除点"""
        if self.current_curve is None or len(self.current_curve) == 0:
            return
            
        # 找到最近的点
        distances = np.sqrt((self.current_curve[:, 0] - self.f[freq_idx])**2 + 
                           (self.current_curve[:, 1] - vs)**2)
        nearest_idx = np.argmin(distances)
        
        # 获取要删除的点的信息
        removed_point = self.current_curve[nearest_idx]
        freq_val, vs_val = removed_point
        
        # 删除最近的点
        self.current_curve = np.delete(self.current_curve, nearest_idx, axis=0)
        
        self.log_action("Point removed", f"Curve {self.current_curve_index}, f={freq_val:.2f} Hz, v={vs_val:.0f} m/s")
        
        # 更新存储
        if len(self.current_curve) == 0:
            self.current_curve = None
            if self.current_curve_index in self.curve_lines:
                try:
                    self.curve_lines[self.current_curve_index].remove()
                except NotImplementedError:
                    # 如果无法直接移除，将其设为不可见
                    self.curve_lines[self.current_curve_index].set_visible(False)
                # 清除标注
                for ann in self.point_annotations.get(self.current_curve_index, []):
                    try:
                        ann.remove()
                    except NotImplementedError:
                        # 如果无法直接移除，将其设为不可见
                        ann.set_visible(False)
                if self.current_curve_index in self.point_annotations:
                    del self.point_annotations[self.current_curve_index]
                del self.curve_lines[self.current_curve_index]
        else:
            self.curves[self.current_curve_index] = self.current_curve
            
        self.draw_curve()
        
    def on_key_press(self, event):
        """处理键盘按键事件"""
        if event.key == 'a':
            self.set_add_mode(None)
        elif event.key == 'm':
            self.set_move_mode(None)
        elif event.key == 'r':
            self.set_remove_mode(None)
        elif event.key == 's':
            self.save_curves(None)
        elif event.key == 'n':
            self.new_curve(None)
        elif event.key == 'd':
            self.delete_current_curve(None)
        elif event.key == 'l':
            self.load_curves_from_file(None)
        elif event.key.isdigit():
            curve_idx = int(event.key)
            if curve_idx in self.curves:
                self.switch_to_curve(curve_idx)
        elif event.key == 'delete':
            # 处理删除键
            print("Delete key pressed - select a point to delete")
                
    def set_add_mode(self, event):
        """设置为添加模式"""
        self.mode = 'add'
        self.log_action("Mode changed", "ADD - Click to add points to curve")
        self.update_status_text()
        
    def set_move_mode(self, event):
        """设置为移动模式"""
        self.mode = 'move'
        self.log_action("Mode changed", "MOVE - Click to move selected point")
        self.update_status_text()
        
    def set_remove_mode(self, event):
        """设置为删除模式"""
        self.mode = 'remove'
        self.log_action("Mode changed", "REMOVE - Right-click to remove points")
        self.update_status_text()
        
    def new_curve(self, event):
        """创建新曲线"""
        new_index = max(self.curves.keys()) + 1 if self.curves else 1
        self.curves[new_index] = np.empty((0, 2))
        self.switch_to_curve(new_index)
        self.log_action("New curve created", f"Curve {new_index}")
        
    def delete_current_curve(self, event):
        """删除当前曲线"""
        deleted_curve_index = self.current_curve_index
        if self.current_curve_index in self.curves:
            # 删除曲线的可视化元素
            if self.current_curve_index in self.curve_lines:
                try:
                    self.curve_lines[self.current_curve_index].remove()
                except NotImplementedError:
                    # 如果无法直接移除，将其设为不可见
                    self.curve_lines[self.current_curve_index].set_visible(False)
                # 清除标注
                for ann in self.point_annotations.get(self.current_curve_index, []):
                    try:
                        ann.remove()
                    except NotImplementedError:
                        # 如果无法直接移除，将其设为不可见
                        ann.set_visible(False)
                if self.current_curve_index in self.point_annotations:
                    del self.point_annotations[self.current_curve_index]
                del self.curve_lines[self.current_curve_index]
            
            # 删除曲线数据
            del self.curves[self.current_curve_index]
            self.current_curve = None
            
            # 切换到其他曲线（如果有的话）
            if self.curves:
                next_index = list(self.curves.keys())[0]
                self.switch_to_curve(next_index)
            else:
                # 如果没有其他曲线，创建一个新的
                self.curves[0] = np.empty((0, 2))
                self.switch_to_curve(0)
                
            self.log_action("Curve deleted", f"Curve {deleted_curve_index}")
            self.fig.canvas.draw()
        
    def switch_to_curve(self, curve_index):
        """切换到指定曲线"""
        self.current_curve_index = curve_index
        self.current_curve = self.curves.get(curve_index, np.empty((0, 2)))
        self.draw_curve()
        self.update_status_text()
        self.log_action("Switched to curve", f"Curve {curve_index}")
        
    def save_curves(self, event):
        """保存拾取的曲线"""
        # 保存到与h5文件相同的路径，但扩展名改为.disp
        base_name = os.path.splitext(self.h5_file)[0]
        output_file = base_name + '.disp'
        
        # 检查文件是否已存在，如果存在则添加时间后缀
        if os.path.exists(output_file):
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = base_name + '_' + timestamp + '.disp'
        
        with open(output_file, 'w') as f:
            f.write("# Picked dispersion curves\n")
            f.write("# Format: frequency velocity\n")
            for curve_idx, curve in self.curves.items():
                if len(curve) > 0:
                    f.write(f"\n# Curve {curve_idx}\n")
                    for freq, vs in curve:
                        f.write(f"{freq:.6f} {vs:.6f}\n")
                        
        self.log_action("Curves saved", f"Saved to {output_file}")
        
    def load_dispersion_curves(self, disp_file):
        """加载已保存的频散曲线"""
        curves = {}
        current_curve_index = None
        current_curve_points = []
        
        with open(disp_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith('# Curve'):
                    # 保存之前的曲线
                    if current_curve_index is not None and current_curve_points:
                        curves[current_curve_index] = np.array(current_curve_points)
                    
                    # 开始新的曲线
                    current_curve_index = int(line.split()[2])
                    current_curve_points = []
                elif line and not line.startswith('#'):
                    # 读取频率和速度数据
                    freq, vs = map(float, line.split())
                    current_curve_points.append([freq, vs])
            
            # 保存最后一条曲线
            if current_curve_index is not None and current_curve_points:
                curves[current_curve_index] = np.array(current_curve_points)
        
        # 更新曲线数据
        self.curves = curves
        if self.curves:
            # 设置当前曲线为第一条曲线
            first_curve_index = min(self.curves.keys())
            self.current_curve_index = first_curve_index
            self.current_curve = self.curves[first_curve_index]
        else:
            self.current_curve_index = 0
            self.current_curve = None
            
        # 绘制所有曲线
        for curve_index in sorted(self.curves.keys()):
            self.current_curve_index = curve_index
            self.current_curve = self.curves[curve_index]
            self.draw_curve()
        
        # 恢复到第一条曲线
        if self.curves:
            first_curve_index = min(self.curves.keys())
            self.switch_to_curve(first_curve_index)

    def load_curves_from_file(self, event):
        """从文件加载曲线"""
        disp_file = os.path.splitext(self.h5_file)[0] + '.disp'
        if os.path.exists(disp_file):
            # 清除现有曲线
            self.clear_all_curves()
            
            # 加载曲线
            self.load_dispersion_curves(disp_file)
            self.log_action("Curves loaded", f"From {disp_file}")
        else:
            self.log_action("Load failed", f"File {disp_file} not found")
        
    def load_reference_curves(self, text):
        """加载参考曲线文件"""
        if not text:
            return
            
        # 检查文件是否存在
        if os.path.exists(text):
            # 清除现有曲线
            self.clear_all_curves()
            
            # 加载曲线
            self.load_dispersion_curves(text)
            self.log_action("Reference curves loaded", f"From {text}")
        else:
            self.log_action("Load failed", f"Reference file {text} not found")
        
    def create_log_file(self):
        """创建日志文件"""
        base_name = os.path.splitext(self.h5_file)[0]
        log_file = base_name + '.log'
        
        # 如果日志文件已存在，则添加时间戳
        if os.path.exists(log_file):
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            log_file = base_name + '_' + timestamp + '.log'
            
        # 创建日志文件并写入头部信息
        with open(log_file, 'w') as f:
            f.write(f"Dispersion Curve Picker Log\n")
            f.write(f"File: {self.h5_file}\n")
            f.write(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("=" * 50 + "\n")
            
        return log_file
    
    def log_action(self, action, details=""):
        """记录操作到日志文件"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        log_entry = f"[{timestamp}] {action}"
        if details:
            log_entry += f" - {details}"
        log_entry += "\n"
        
        # 打印到控制台
        print(log_entry.strip())
        
        # 写入日志文件
        with open(self.log_file, 'a') as f:
            f.write(log_entry)
            
    def refresh_h5_file(self, event):
        """刷新H5文件并重新初始化"""
        h5_file = self.txt_h5.text.strip()
        if not h5_file or not os.path.exists(h5_file):
            self.log_action("Refresh failed", f"H5 file {h5_file} not found")
            return
            
        # 保存当前曲线
        old_curves = self.curves.copy()
        
        # 重新初始化
        self.__init__(h5_file)
        
        # 恢复曲线
        self.curves = old_curves
        for curve_index in sorted(self.curves.keys()):
            self.current_curve_index = curve_index
            self.current_curve = self.curves[curve_index]
            self.draw_curve()
            
        # 切换到第一条曲线
        if self.curves:
            first_curve_index = min(self.curves.keys())
            self.switch_to_curve(first_curve_index)
            
        self.log_action("H5 file refreshed", f"Loaded {h5_file}")
        
    def show(self):
        """显示图形界面"""
        plt.show()


def main():
    parser = argparse.ArgumentParser(description='Interactive dispersion curve picker')
    parser.add_argument('h5_file', help='Input H5 file containing FV data')
    parser.add_argument('--disp-file', help='Reference dispersion curve file (.disp)')
    args = parser.parse_args()
    
    picker = DispersionCurvePicker(args.h5_file)
    
    # 如果提供了参考disp文件，则加载它
    if args.disp_file and os.path.exists(args.disp_file):
        picker.reference_file = args.disp_file
        picker.clear_all_curves()
        picker.load_dispersion_curves(args.disp_file)
        picker.log_action("Reference curves loaded", f"From {args.disp_file}")
        
    picker.show()


if __name__ == '__main__':
    main()